import os
import json
import datetime
from functools import partial

import numpy as np
import matplotlib.pyplot as plt
import torch


class SampleLogger:

    sep = "-"


    def __init__(self, slogger_dir="", **kwargs):

        if slogger_dir == "":
            self.save_dir = None
            return

        now = datetime.datetime.now()
        folder = now.strftime("%Y-%m-%d-&-%H-%M-%S")

        self.save_dir = os.path.join(slogger_dir, folder)
        os.makedirs(self.save_dir, exist_ok=True)

        self.json_saver = JsonSaver(self.save_dir)
        self.tensor_saver = TensorSaver(self.save_dir)

        self.fn_dict = {}
        for saver in [self.json_saver, self.tensor_saver]:
            self.fn_dict.update(saver.pattern)


    def info(self, num=0, **kwargs):

        if self.save_dir is None:
            return

        for k, v in kwargs.items():
            split_ = k.split(self.sep)
            attrname = split_[0]
            pattern = split_[1]

            fn = self.fn_dict[pattern]
            fn(v, name=attrname, num=num)


class JsonSaver:
    """
    filename : sclar name - iter num - num in batch
    """

    def __init__(self, save_dir=None):

        self.save_dir = os.path.join(save_dir, "json")

        os.makedirs(
            self.save_dir, exist_ok=True
        )

        self.pattern = {
            "scalar" : partial(self.save, name="scalar")
        }


    def save(self, value, *, name="", num=0, **kwargs):
        
        log = {}

        for k, v in value.items():

            suffix = "-".join([str(k), str(num)])
            for i, s in enumerate(v):
                
                keyname = "-".join([suffix, str(i+1)])

                if type(s) == torch.Tensor:
                    if len(s.shape) == 0:
                        s = s.item()
                    else:
                        s = s.tolist()
                log[keyname] = s

        json_name = os.path.join(
            self.save_dir, "-".join([name, str(num)])+".json"
        )

        with open(json_name, "w") as fp:
            json.dump(log, fp, indent='\t', sort_keys=True)


class TensorSaver:
    """
    filename : data name - iter num - num in batch - suffix
    """

    def __init__(self, save_dir=None):

        self.save_dir = save_dir

        self.pattern = {
            "img" : partial(self.save, folder="img", transform=self.img_trans, cmap="gray"),
            "fmap_chs" : partial(self.save, folder="fmap_chs", transform=self.fm_ch1)
        }

        for k in self.pattern.keys():
            os.makedirs(
                os.path.join(self.save_dir, k), exist_ok=True
            )


    def save(
        self, value, *, name="", num=0, folder="", 
        transform=None, **kwargs
    ):
        cmap = "viridis" 
        if "cmap" in kwargs:
            cmap = kwargs["cmap"]

        if not transform is None:
            ret = {}
            ret.update(transform(value))
        else:
            ret = {"" : value}

        for k, v in ret.items():

            # if len(v.shape) < 4:
            #     v = np.expand_dims(v, 0)

            cnt = 0

            for image in v:
                cnt += 1
                name_seq = [name, str(num)]
                name_seq.append(str(cnt))

                if k != "":
                    name_seq.append(k)

                filename = os.path.join(
                    self.save_dir, folder, "-".join(name_seq)+".png"
                )
                
                image = image.squeeze()

                plt.imsave(filename, image, cmap=cmap)


    def img_trans(self, value):

        if value.device != "cpu":
            value = value.cpu()

        value = value.numpy()

        if len(value.shape) < 4:
            value = np.expand_dims(value, 1)

        value = np.transpose(value, [0, 2, 3, 1])
        value = (value - value.min())/(value.ptp())

        return {"" : value}


    def fm_ch1(self, value):

        if value.device != "cpu":
            value = value.cpu()

        value = value.numpy()

        if len(value.shape) < 4:
            value = np.expand_dims(value, 1)

        value = np.transpose(value, [0, 2, 3, 1])
        value = (value - value.min())/(value.ptp())
        channels = value.shape[3]

        ret = {}
        for i in range(channels):
            ret["ch"+str(i+1)] = np.expand_dims(value[:, :, :, i], 3)

        return ret


if __name__ == "__main__1":

    class test:

        a = 0

        def __init__(self):
            self.p = {
                "a" : self.print
            }

        def add(self):
            self.a += 1

        def print(self):
            print(self.a)
            print("----------")


    t = test()

    fn = t.p["a"]
    tmp = fn()

    t.add()

    tmp = fn()


if __name__ == "__main__1":
    import torch

    save_dir = "E:/datasets/client-data"

    ts = TensorSaver(save_dir=save_dir)

    # fn = ts.pattern["img"]
    fn = ts.pattern["fmap_chs"]

    tensor = torch.rand((4, 10, 256, 256))

    fn(tensor, name="fuck", num=9527)


if __name__ == "__main__1":
    import torch

    save_dir = "E:/datasets/client-data"

    js = JsonSaver(save_dir=save_dir)

    fn = js.pattern["scalar"]

    scalar = {
        "A" : torch.randint(0, 10, (3, )), 
        "B" : [1, 3, 4],
        "c" : [(1, 3), (3, 4), (5, 6)] 
    }

    fn(scalar, name="fuck", num=9527)


if __name__ == "__main__":
    import torch

    save_dir = "E:/datasets/client-data"

    sl = SampleLogger(slogger_dir=save_dir)

    
    sar_img = torch.randint(0, 127, (4, 3, 256, 256)) 
    opt_img = torch.randint(127, 255, (4, 3, 256, 256))
    
    sar_ftm = torch.rand((4, 6, 256, 256)) 
    opt_ftm = torch.rand((4, 6, 256, 256))

    scalar = {
        "dis" : torch.randint(0, 10, (4, )), 
        "B" : [1, 3, 4, 8],
        "c" : [(1, 3), (3, 4), (5, 6), (10, 11)] 
    }

    input = {
        "sar-img" : sar_img, 
        "opt-img" : opt_img, 
        "sar-fmap_chs" : sar_ftm, 
        "opt-fmap_chs" : opt_ftm, 
        "scalar-scalar" : scalar
    }

    sl.info(num=666, **input)


